package com.practice.niuke.new_direct_practice.class27;

import java.util.HashMap;
import java.util.List;
import java.util.Stack;

/**
 * 并查集的经典实现
 */
public class UnionFind {

    // 样本进来会包一层，叫做元素
    public static class Element<V> {
        public V value;

        public Element(V value) {
            this.value = value;
        }

    }

    // 并查集，V代表并查集中存储的元素的类型
    public static class UnionFindSet<V> {

        // 样本类型，实例 -> element
        // a -> a的点
        public HashMap<V, Element<V>> elementMap;
        // key  某个元素  value 该元素的父
        public HashMap<Element<V>, Element<V>> fatherMap;
        // key 某个集合的代表元素   value 该集合的大小
        // 如果某个Element（某个点），不是集合的代表点，sizemap里没有它的记录
        // 如果某个Element（某个点），是集合的代表点，sizemap里才会有它的记录，记录的值表示这个点代表的集合的大小
        public HashMap<Element<V>, Integer> sizeMap;

        public UnionFindSet(List<V> list) {
            elementMap = new HashMap<>();
            fatherMap = new HashMap<>();
            sizeMap = new HashMap<>();

            for (V value : list) {
                // 初始化的时候对所有的样本数据，进行并查集节点的构造，
                Element<V> element = new Element<V>(value);
                elementMap.put(value, element);
                // 每个节点单独作为一个集合，这个节点就是其集合的代表节点（头节点），那么它的父节点就是它自己
                fatherMap.put(element, element);
                // 初始化是每个新生成的集合的节点数为1
                sizeMap.put(element, 1);
            }
        }

        // 给定一个ele，往上一直找，把这个ele所在集合的代表元素（头元素）返回
        private Element<V> findHead(Element<V> element) {
            // 沿途所有路过的节点全记录
            Stack<Element<V>> path = new Stack<>();
            while (element != fatherMap.get(element)) {
                path.push(element);
                element = fatherMap.get(element);
            }
            // 把路径中所有的节点都弹出，直接挂在集合的代表节点（头节点）下，也就是把路径中的节点的父变为代表节点
            while (!path.isEmpty()) {
                fatherMap.put(path.pop(), element);
            }
            return element;
        }

        /**
         * 查询a所在的集合和b所在的集合是不是同一个集合
         *
         * @param a a
         * @param b b
         * @return boolean
         */
        public boolean isSameSet(V a, V b) {
            if (elementMap.containsKey(a) && elementMap.containsKey(b)) {
                return findHead(elementMap.get(a)) == findHead(elementMap.get(b));
            }
            return false;
        }

        /**
         * 合并a所在的集合和b所在的集合
         *
         * @param a a
         * @param b b
         */
        public void union(V a, V b) {
            if (elementMap.containsKey(a) && elementMap.containsKey(b)) {
                Element<V> aF = findHead(elementMap.get(a));
                Element<V> bF = findHead(elementMap.get(b));
                if (aF != bF) {
                    // 谁所在的集合大
                    Element<V> big = sizeMap.get(aF) >= sizeMap.get(bF) ? aF : bF;
                    // 谁所在的集合小
                    Element<V> small = big == aF ? bF : aF;
                    // 更新 fatherMap
                    fatherMap.put(small, big);
                    // 更新 sizeMap
                    sizeMap.put(big, sizeMap.get(aF) + sizeMap.get(bF));
                    // 移除sizeMap中小集合的代表点
                    sizeMap.remove(small);
                }
            }
        }

    }

}
